import os
import subprocess

class AudioExtractor:
    """
    批量提取视频音频工具（支持自定义输出格式与参数）
    """

    def __init__(self, input_folder, output_folder=r"C:\Users\33720\Desktop", ffmpeg_path="ffmpeg"):
        self.input_folder = input_folder
        self.output_folder = output_folder
        self.ffmpeg_path = ffmpeg_path

        # 默认参数
        self.output_format = "mp3"
        self.codec = "libmp3lame"
        self.sample_rate = 44100
        self.channels = 2
        self.bitrate = "192k"

        os.makedirs(self.output_folder, exist_ok=True)

    # -------------------- Setter 方法 --------------------
    def set_output_format(self, fmt: str):
        """设置输出音频格式（如 mp3, wav, aac, flac 等）"""
        self.output_format = fmt.lower()
        # 根据常见格式自动调整默认编码器
        codec_map = {
            "mp3": "libmp3lame",
            "aac": "aac",
            "wav": "pcm_s16le",
            "flac": "flac",
            "ogg": "libvorbis"
        }
        self.codec = codec_map.get(self.output_format, "copy")
        return self

    def set_bitrate(self, bitrate: str):
        """设置音频比特率，例如 '192k'、'256k'"""
        self.bitrate = bitrate
        return self

    def set_sample_rate(self, rate: int):
        """设置采样率，例如 44100"""
        self.sample_rate = rate
        return self

    def set_channels(self, channels: int):
        """设置声道数，例如 2（立体声）"""
        self.channels = channels
        return self

    def set_codec(self, codec: str):
        """手动设置编码器"""
        self.codec = codec
        return self

    def set_output_folder(self, output_folder: str):
            """手动设置编码器"""
            self.output_folder = output_folder
            return self

    # -------------------- 主功能 --------------------
    def extract_audio(self):
        os.makedirs(self.output_folder, exist_ok=True)
        """批量提取音频"""
        video_exts = (".mp4", ".mkv", ".mov", ".avi", ".flv", ".ts")
        files = [f for f in os.listdir(self.input_folder)
                 if f.lower().endswith(video_exts)]

        if not files:
            print("⚠️ 未找到任何支持的视频文件。")
            return
        result=[]
        for filename in files:
            input_path = os.path.join(self.input_folder, filename)
            base_name = os.path.splitext(filename)[0]
            output_path = os.path.join(self.output_folder, f"{base_name}.{self.output_format}")
            result.append(output_path)
            print(f"🎬 正在提取音频：{filename} → {base_name}.{self.output_format}")

            cmd = [
                self.ffmpeg_path,
                "-i", input_path,
                "-vn",                    # 去除视频流
                "-acodec", self.codec,
                "-ar", str(self.sample_rate),
                "-ac", str(self.channels),
                "-b:a", self.bitrate,
                "-y",                     # 覆盖输出文件
                output_path
            ]

            try:
                subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
                print(f"✅ 成功提取：{output_path}")
            except subprocess.CalledProcessError as e:
                print(f"❌ 提取失败：{filename}\n错误信息：{e.stderr.decode(errors='ignore')}")

        print("\n🎉 所有文件已处理完成！")
        return result
# -------------------- 示例使用 --------------------
if __name__ == "__main__":
    input_folder = input("📂 请输入视频文件夹路径：").strip('"')
    output_folder = input("💾 请输入输出音频文件夹路径：").strip('"')

    extractor = AudioExtractor(input_folder, output_folder)

    # 可灵活配置
    extractor.set_output_format("flac") \
             .set_bitrate("256k") \
             .set_sample_rate(48000) \
             .set_channels(2)

    extractor.extract_audio()
